package oneD.fem.method;

import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.GeneralMatrixLSESolver;
import inf.jlinalg.ILSESolver;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IVector;
import inf.jlinalg.QuadraticMatrixInfo;
import inf.jlinalg.SolveFailedException;
import inf.jlinalg.Vector3D;

import java.util.ArrayList;

import oneD.fem.algorithm.MyLinalg;
import oneD.fem.structure.Structure;

public class ArcLengthControl {

	public class LoadDisplacementState {
		Vector3D[] displacement_;
		double load_;

		public LoadDisplacementState(Structure struct, double lambda,
				IVector uGlobal) {
			this.displacement_ = new Vector3D[struct.getNumberOfNodes()];
			this.load_ = lambda;
			this.setDisplacement(struct, uGlobal);
		}

		public double getLoad() {
			return this.load_;
		}

		public Vector3D getDisplacement(int i) {
			return this.displacement_[i];
		}

		private void setDisplacement(Structure struct, IVector uGlobal) {
			for (int i = 0; i < struct.getNumberOfNodes(); i++) {
				double[] u = new double[3];
				int[] dofNumbers = struct.getNode(i).getDOFNumbers();
				for (int j = 0; j < 3; j++)
					if (dofNumbers[j] == -1)
						u[j] = 0;
					else
						u[j] = uGlobal.get(dofNumbers[j]);
				this.displacement_[i] = new Vector3D(u);
			}
		}

		public void print() {
			System.out.println("lambda = " + this.getLoad());
			for (int i = 0; i < this.displacement_.length; i++)
				System.out.println("u" + i + "="
						+ this.getDisplacement(i).toString());
		}
	}

	private static final int MAX_ITERATION = 1000;
	private static final double TOL = 1e-6;

	private ArrayList<LoadDisplacementState> curve_;
	private Structure struct_;
	private int constraintType_;

	private double s_; // constraint parameter
	private double psi_ = 1;
	private int N_;
	private IVector u0_;
	private double lambda0_;

	public ArcLengthControl(Structure struct) {
		this.curve_ = new ArrayList<LoadDisplacementState>();
		this.struct_ = struct;
	}

	public void clear() {
		/*
		 * be careful with this method. After this if call get.. method will
		 * cause error
		 */
		this.curve_.clear();
	}

	public void setConstraintType(String type) {
		if (type == "load control")
			this.constraintType_ = 1;
		else if (type == "displacement control")
			this.constraintType_ = 2;
		else if (type == "sphere")
			this.constraintType_ = 3;
		else
			this.constraintType_ = 0;
	}

	public void setNumStep(int n) {
		this.N_ = n;
	}

	public int getNumStep() {
		return this.N_;
	}

	public Structure getStructure() {
		return this.struct_;
	}


	public void setPsi(double psi) {
		this.psi_ = psi;
	}

	public double getPsi() {
		return this.psi_;
	}

	public void setParameterS(double s) {
		this.s_ = s;
	}

	public double getParameterS() {
		return this.s_;
	}

	public void setU0(IVector u0) {
		this.u0_ = u0;
	}

	public void setLambda0(double lambda0) {
		this.lambda0_ = lambda0;
	}

	public double getLoad(int step) {
		return this.curve_.get(step).getLoad();
	}

	public Vector3D getDisplacement(int step, int node) {
		return this.curve_.get(step).getDisplacement(node);
	}

	private double f(IVector u, IVector un, double lambda, double lambdan,
			double psi, double s) {
		switch (this.constraintType_) {
		case 1:
			return this.loadControlConstraint(lambda, s);
		case 2:
			return this.displacementControlConstraint(u, s);
		case 3:
			return this.sphereSurfaceControl(u, un, lambda, lambdan, psi, s);
		default:
			return 0;
		}
	}

	private IVector dfdu(IVector u, IVector un, double lambda, double lambdan,
			double psi, double s) {
		switch (this.constraintType_) {
		case 1:
			return this.loadControlConstraintDerivativeWrtU(u);
		case 2:
			return this.displacementControlConstraintDerivativeWrtU(u);
		case 3:
			return this.sphereSurfaceControlDerivativeWrtU(u, un, lambda,
					lambdan, psi, s);
		default:
			return null;
		}
	}

	private double dfdlambda(IVector u, IVector un, double lambda,
			double lambdan, double psi, double s) {
		switch (this.constraintType_) {
		case 1:
			return this.loadControlConstraintDerivativeWrtLambda();
		case 2:
			return this.displacementControlConstraintDerivativeWrtLambda();
		case 3:
			return this.sphereSurfaceControlDerivativeWrtLambda(u, un, lambda,
					lambdan, psi, s);
		default:
			return 0;
		}
	}

	public void run() {
		int NEQ = this.struct_.enumerateDOFs();
		// solver initialization
		ILSESolver solver = new GeneralMatrixLSESolver(
				GeneralMatrixLSESolver.SOLVER_IMPL_JAVA);
		// using this for solver implementation in Java, it will be slow
		QuadraticMatrixInfo aInfo = solver.getAInfo();
		IMatrix ktGlobal = solver.getA();
		IVector rGlobal = new ArrayVector(NEQ);
		IVector riGlobal = new ArrayVector(NEQ);
		IVector delta_u_lambda;
		IVector delta_u_r;
		double delta_lambda;

		aInfo.setSize(NEQ);

		struct_.assembleLoadVector(rGlobal);

		IVector u = (IVector) u0_.clone();
		double lambda = lambda0_;

		int cnt = 0;
		boolean cond = false;
		do {
			cnt++;
			// store the result first
			this.curve_.add(new LoadDisplacementState(this.struct_, lambda, u));

			IVector um = (IVector) u.clone();
			double lambdam = lambda;

			/*
			 * predictor step
			 */
			solver.initialize();
			System.out.println("Analysis tracking step: " + cnt);

			BLAM.zero(ktGlobal);// clear ktGlobal to fill in new data

			this.struct_.setDisplacement(u);
			this.struct_.assembleTangentStiffnessMatrix(ktGlobal);

			double[] b = MyLinalg.toArray(rGlobal);
			try {
				solver.solve(b);
			} catch (SolveFailedException e) {
				System.out.println("Solver failed: " + e.getMessage());
			}

			delta_u_lambda = new ArrayVector(b);
			double s0 = Math.sqrt(1 + Math.pow(
					BLAM.norm(BLAM.NORM_TWO, delta_u_lambda), 2));
			lambda = lambda + s_ / s0;
			MyLinalg.scale(s_ / s0, delta_u_lambda);
			MyLinalg.add(1, delta_u_lambda, u);

			/*
			 * corrector step
			 */
			double epsilon = 1.0;
			int i_cnt = 0;
			do {
				this.struct_.setDisplacement(u);
				solver.initialize();
				BLAM.zero(ktGlobal);
				this.struct_.assembleTangentStiffnessMatrix(ktGlobal);
				BLAM.fill(riGlobal, 0.0);
				this.struct_.assembleInternalForceVector(riGlobal);

				MyLinalg.add(-lambda, rGlobal, riGlobal);
				MyLinalg.scale(-1, riGlobal);

				b = MyLinalg.toArray(riGlobal);
				try {
					solver.solve(b);
				} catch (SolveFailedException e) {
					System.out.println("Solver failed: " + e.getMessage());
				}
				delta_u_r = new ArrayVector(b);

				b = MyLinalg.toArray(rGlobal);
				try {
					solver.solve(b);
				} catch (SolveFailedException e) {
					System.out.println("Solver failed: " + e.getMessage());
				}
				delta_u_lambda = new ArrayVector(b);

				IVector dfdu_ = this.dfdu(u, um, lambda, lambdam, psi_, s_);
				double f_ = this.f(u, um, lambda, lambdam, psi_, s_);
				double dfdlambda_ = this.dfdlambda(u, um, lambda, lambdam,
						psi_, s_);
				delta_lambda = -(f_ + BLAM.dot(dfdu_, delta_u_r))
						/ (BLAM.dot(dfdu_, delta_u_lambda) + dfdlambda_);

				IVector delta_u = (IVector) delta_u_r.clone();
				MyLinalg.add(delta_lambda, delta_u_lambda, delta_u);

				// update
				MyLinalg.add(1, delta_u, u);
				lambda += delta_lambda;

				double norm_delta_u = BLAM.norm(BLAM.NORM_TWO, delta_u);
				IVector temp = (IVector) u.clone();
				MyLinalg.add(-1, um, temp);
				epsilon = BLAM.norm(BLAM.NORM_TWO, temp);
				double epsilon_new;
				if (epsilon < 1e-12)
					epsilon_new = norm_delta_u;
				else
					epsilon_new = norm_delta_u / epsilon;
				i_cnt++;
				if (epsilon_new > epsilon)
					System.out.print("/");
				else
					System.out.print("\\");
				epsilon = epsilon_new;
			} while (epsilon > TOL && i_cnt < MAX_ITERATION);
			if (i_cnt >= MAX_ITERATION)
				System.out.println("The corrector does not converge at step "
						+ cnt + " within " + ArcLengthControl.MAX_ITERATION
						+ " corrector steps");
			System.out.println("error = " + epsilon);
			if (this.constraintType_ == 1 || this.constraintType_ == 2)
				cond = (epsilon > TOL);
			else if (this.constraintType_ == 3)
				cond = (cnt < N_);
		} while (cond);
		this.curve_.add(new LoadDisplacementState(this.struct_, lambda, u));
		// System.out.println("lambda = " + lambda);
		// System.out.println("u = " + MatrixFormat.format(u));
	}

	public void print() {
		for (int i = 0; i < this.curve_.size(); i++) {
			this.curve_.get(i).print();
		}
	}

	public Plot1D plot() {
		Plot1D p = new Plot1D(this);
		return p;
	}

	private double loadControlConstraint(double lambda, double s) {
		return lambda - s;
	}

	private IVector loadControlConstraintDerivativeWrtU(IVector u) {
		IVector dfdu = new ArrayVector(u.getSize());
		BLAM.fill(dfdu, 0.0);
		return dfdu;
	}

	private double loadControlConstraintDerivativeWrtLambda() {
		return 1;
	}

	private double displacementControlConstraint(IVector u, double s) {
		return BLAM.norm(BLAM.NORM_TWO, u) - s;
	}

	private IVector displacementControlConstraintDerivativeWrtU(IVector u) {
		IVector du = (IVector) u.clone();
		MyLinalg.scale(1 / BLAM.norm(BLAM.NORM_TWO, u), du);
		return du;
	}

	private double displacementControlConstraintDerivativeWrtLambda() {
		return 0;
	}

	private double sphereSurfaceControl(IVector u, IVector un, double lambda,
			double lambdan, double psi, double s) {
		IVector temp = (IVector) u.clone();
		MyLinalg.add(-1, un, temp);
		double f = Math.pow(BLAM.norm(BLAM.NORM_TWO, temp), 2);
		f += psi * psi * (lambda - lambdan) * (lambda - lambdan);
		return Math.sqrt(f) - s;
	}

	private IVector sphereSurfaceControlDerivativeWrtU(IVector u, IVector un,
			double lambda, double lambdan, double psi, double s) {
		IVector du = (IVector) u.clone();
		MyLinalg.add(-1, un, du);
		double f = this.sphereSurfaceControl(u, un, lambda, lambdan, psi, s)
				+ s;
		MyLinalg.scale(1 / f, du);
		return du;
	}

	private double sphereSurfaceControlDerivativeWrtLambda(IVector u,
			IVector un, double lambda, double lambdan, double psi, double s) {
		double f = this.sphereSurfaceControl(u, un, lambda, lambdan, psi, s)
				+ s;
		return psi * psi * (lambda - lambdan) / f;
	}
}
